TopKFusion

沿指定轴查找最大或最小 topk 个值的索引。当 topk=1 且 largest=true 时，该算子等价于 ArgMax；当 topk=1 且 largest=false 时，该算子等价于 ArgMin。

对于输入张量沿指定轴的每个切片 \(X_{slice_i}\)，本算子返回其中最大（largest=true）或最小（largest=false）的 topk 个元素的索引：

\[Y_{i,j} = \operatorname{TopKIndex}_{j}(X_{slice_i}, largest), \quad j = 0, 1, \dots, topk-1\]

其中 \(Y_{i,j}\) 表示第 \(j\) 个结果的索引，返回结果按值从大到小（largest=true）或从小到大（largest=false）排列。

输入：

  • input - 输入数据地址。

  • params - 其他参数打包成数组。

  • core_mask - 核掩码（仅共享存储版本使用）。

输出：

  • output - 存储索引的输出张量，数据类型为 int32。

  • output_value - 存储找到的值，数据类型与输入相同。

备注

每次调用都会同时填充 output`（索引）和 `output_value`（值），不存在 `return_values 开关。

备注

返回的 topk 个结果按值有序排列： - 当 largest=true 时，按值从大到小排列； - 当 largest=false 时，按值从小到大排列。

支持平台：

FT78NE MT7004

备注

  • FT78NE 支持 fp32、fp64、int32、int16、int8

  • MT7004 支持 fp16、fp32、int32、int16

参数数组结构：

 1long long params[11];
 2params[0]  = (long long)in_shape;      // 输入张量的维度信息数组
 3params[1]  = (long long)in_strides;    // 输入张量的步长信息数组
 4params[2]  = (long long)out_strides;   // 输出张量的步长信息数组
 5params[3]  = (long long)arg_elements;  // 用于存放候选值的临时工作空间地址，大小为 axis_dim * core_num * sizeof(输入数据类型)，8 字节对齐
 6params[4]  = (long long)index;         // 用于存放候选索引的临时工作空间地址，大小为 axis_dim * core_num * sizeof(int32)，8 字节对齐
 7params[5]  = (long long)topk;          // 需要查找的最大/最小值的数量
 8params[6]  = (long long)in_shape_size; // 输入张量的维度数 (即 in_shape 数组的长度)
 9params[7]  = (long long)axis;          // 执行查找操作的轴
10params[8]  = (long long)largest;       // 是否查找最大值的标志。若为 0，则查找最小值
11params[9]  = (long long)topk_val;      // 临时缓存，大小为 topk * core_num * sizeof(输入数据类型)，8 字节对齐
12params[10] = (long long)topk_idx;      // 临时缓存，大小为 topk * core_num * sizeof(int32)，8 字节对齐

共享存储版本:

void fp_topk_s(float *input, void *output, float *output_value, long long *params, int core_mask)
void hp_topk_s(half *input, void *output, half *output_value, long long *params, int core_mask)
void dp_topk_s(double *input, void *output, double *output_value, long long *params, int core_mask)
void i32_topk_s(int *input, void *output, int *output_value, long long *params, int core_mask)
void i16_topk_s(int16_t *input, void *output, int16_t *output_value, long long *params, int core_mask)
void i8_topk_s(int8_t *input, void *output, int8_t *output_value, long long *params, int core_mask)

C调用示例：

 1// MT7004 示例
 2#include <stdio.h>
 3#include <topk.h>
 4int main(int argc, char* argv[]) {
 5    float* input = (float*)0x81000000; // 需要初始化
 6    void* output = (void*)0x82000000;  // 不需要初始化，输出 int32 索引
 7    float* output_value = (float*)0x83000000; // 不需要初始化，输出对应值
 8    float* arg_elements = (float*)0x84000000; // 不需要初始化
 9    float* topk_val = (float *)0x89000000;
10    int* topk_idx = (int *)0x8A000000;
11    int* index = (int*)0x85000000; // 不需要初始化
12
13    int core_mask = 0b1111;
14
15    int *in_strides = (int*)0x86000000;
16    int *out_strides = (int*)0x86000200;
17
18    int in_shape_size = 4; // 最多只考虑4维
19    int in_shape[4] = {4, 8, 16, 8};
20
21    int axis = 1; // 要操作的维度，不能超过3
22    int topk = 3; // 不超过 in_shape[axis]
23
24    srand(time(0));
25    // 初始化测试数据，包含各种情况
26    int i;
27    int in_total_elements = in_shape[0] * in_shape[1] * in_shape[2] * in_shape[3];
28    for(i = 0; i < in_total_elements; i ++) {
29        input[i] = (float)(rand()%100);
30    }
31
32    long long params[11];
33    params[0] = (long long)in_shape;
34    params[1] = (long long)in_strides;
35    params[2] = (long long)out_strides;
36    params[3] = (long long)arg_elements;
37    params[4] = (long long)index;
38    params[5] = (long long)topk;
39    params[6] = (long long)in_shape_size;
40    params[7] = (long long)axis;
41    params[8] = (long long)1; // largest: 1 表示最大值，0 表示最小值
42    params[9] = (long long)topk_val;
43    params[10] = (long long)topk_idx;
44
45    fp_topk_s(input, output, output_value, params, core_mask);
46    return 0;
47}

私有存储版本:

void fp_topk_p(float *input, void *output, float *output_value, long long *params)
void hp_topk_p(half *input, void *output, half *output_value, long long *params)
void dp_topk_p(double *input, void *output, double *output_value, long long *params)
void i32_topk_p(int *input, void *output, int *output_value, long long *params)
void i16_topk_p(int16_t *input, void *output, int16_t *output_value, long long *params)
void i8_topk_p(int8_t *input, void *output, int8_t *output_value, long long *params)

C调用示例：

 1// MT7004 示例
 2#include <stdio.h>
 3#include <topk.h>
 4int main(int argc, char* argv[]) {
 5    float* input = (float*)0x10010000; // 需要初始化
 6    void* output = (void*)0x10020000;  // 不需要初始化，输出 int32 索引
 7    float* output_value = (float*)0x10030000; // 不需要初始化，输出对应值
 8    float* arg_elements = (float*)0x10040000; // 不需要初始化
 9    float* topk_val = (float *)0x10048000;
10    int* topk_idx = (int *)0x10058000;
11    int* index = (int*)0x10050000; // 不需要初始化
12
13    int *in_strides = (int*)0x1004E000;
14    int *out_strides = (int*)0x1004E200;
15
16    int in_shape_size = 4; // 最多只考虑4维
17    int in_shape[4] = {4, 8, 16, 8};
18
19    int axis = 1; // 要操作的维度，不能超过3
20    int topk = 3; // 不超过 in_shape[axis]
21
22    srand(time(0));
23    // 初始化测试数据，包含各种情况
24    int i;
25    int in_total_elements = in_shape[0] * in_shape[1] * in_shape[2] * in_shape[3];
26    for(i = 0; i < in_total_elements; i ++) {
27        input[i] = (float)(rand()%100);
28    }
29
30    long long params[11];
31    params[0] = (long long)in_shape;
32    params[1] = (long long)in_strides;
33    params[2] = (long long)out_strides;
34    params[3] = (long long)arg_elements;
35    params[4] = (long long)index;
36    params[5] = (long long)topk;
37    params[6] = (long long)in_shape_size;
38    params[7] = (long long)axis;
39    params[8] = (long long)1; // largest: 1 表示最大值，0 表示最小值
40    params[9] = (long long)topk_val;
41    params[10] = (long long)topk_idx;
42
43    fp_topk_p(input, output, output_value, params);
44    return 0;
45}